#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库初始化脚本
用于创建PostgreSQL数据库和Redis配置
"""

import sys
import os
import subprocess
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT

# 添加src目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from src.utils import load_config


def check_postgres_connection(host, port, username, password):
    """检查PostgreSQL连接"""
    try:
        # 连接到默认的postgres数据库
        conn = psycopg2.connect(
            host=host,
            port=port,
            user=username,
            password=password,
            database='postgres'
        )
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        return conn
    except Exception as e:
        print(f"❌ PostgreSQL连接失败: {e}")
        return None


def create_database(conn, database_name):
    """创建数据库"""
    try:
        cursor = conn.cursor()
        
        # 检查数据库是否已存在
        cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s", (database_name,))
        exists = cursor.fetchone()
        
        if exists:
            print(f"✅ 数据库 '{database_name}' 已存在")
            return True
        
        # 创建数据库
        cursor.execute(f"CREATE DATABASE {database_name}")
        print(f"✅ 数据库 '{database_name}' 创建成功")
        
        cursor.close()
        return True
        
    except Exception as e:
        print(f"❌ 创建数据库失败: {e}")
        return False


def check_redis_connection(host, port, password=None):
    """检查Redis连接"""
    try:
        import redis
        
        r = redis.Redis(
            host=host,
            port=port,
            password=password,
            decode_responses=True
        )
        
        # 测试连接
        r.ping()
        print(f"✅ Redis连接成功: {host}:{port}")
        return True
        
    except ImportError:
        print("❌ Redis Python包未安装，请运行: pip install redis")
        return False
    except Exception as e:
        print(f"❌ Redis连接失败: {e}")
        return False


def show_installation_guide():
    """显示安装指南"""
    print("\n📋 安装指南:")
    print("=" * 50)
    
    print("\n🐘 PostgreSQL安装:")
    print("Ubuntu/Debian:")
    print("  sudo apt-get update")
    print("  sudo apt-get install postgresql postgresql-contrib")
    print("  sudo systemctl start postgresql")
    print("  sudo systemctl enable postgresql")
    
    print("\n🔴 Redis安装:")
    print("Ubuntu/Debian:")
    print("  sudo apt-get update")
    print("  sudo apt-get install redis-server")
    print("  sudo systemctl start redis-server")
    print("  sudo systemctl enable redis-server")
    
    print("\n🐍 Python依赖安装:")
    print("  pip install -r requirements.txt")
    
    print("\n⚙️  配置说明:")
    print("1. 编辑 config/config.yaml 文件")
    print("2. 设置数据库连接信息")
    print("3. 运行 python init_database.py 初始化数据库")
    print("4. 运行 python iterative_search_demo.py 开始搜索")


def main():
    """主函数"""
    print("🗄️  数据库初始化工具")
    print("=" * 50)
    
    try:
        # 加载配置
        config = load_config()
        db_config = config.get('database', {})
        
        postgres_config = db_config.get('postgres', {})
        redis_config = db_config.get('redis', {})
        
        print("📊 当前配置:")
        print(f"  PostgreSQL: {postgres_config.get('host', 'localhost')}:{postgres_config.get('port', 5432)}")
        print(f"  Redis: {redis_config.get('host', 'localhost')}:{redis_config.get('port', 6379)}")
        print()
        
        # 检查PostgreSQL
        print("🔍 检查PostgreSQL...")
        postgres_conn = check_postgres_connection(
            postgres_config.get('host', 'localhost'),
            postgres_config.get('port', 5432),
            postgres_config.get('username', 'postgres'),
            postgres_config.get('password', '')
        )
        
        if postgres_conn:
            # 创建数据库
            database_name = postgres_config.get('database', 'keyword_mining')
            if create_database(postgres_conn, database_name):
                print("✅ PostgreSQL数据库初始化完成")
            postgres_conn.close()
        else:
            print("❌ PostgreSQL未正确配置或未启动")
        
        # 检查Redis
        print("\n🔍 检查Redis...")
        check_redis_connection(
            redis_config.get('host', 'localhost'),
            redis_config.get('port', 6379),
            redis_config.get('password')
        )
        
        print("\n" + "=" * 50)
        
        # 显示状态
        if postgres_conn:
            print("✅ PostgreSQL: 就绪")
        else:
            print("❌ PostgreSQL: 未就绪")
        
        if check_redis_connection(redis_config.get('host', 'localhost'), redis_config.get('port', 6379)):
            print("✅ Redis: 就绪")
        else:
            print("❌ Redis: 未就绪")
        
        print("\n💡 下一步:")
        if postgres_conn:
            print("  - 运行 python iterative_search_demo.py 开始搜索")
        else:
            print("  - 请先安装并配置PostgreSQL")
            show_installation_guide()
        
    except Exception as e:
        print(f"\n❌ 初始化失败: {e}")
        show_installation_guide()


if __name__ == "__main__":
    main()
